from threading import Thread
import tkinter as tk
from tkinter import filedialog
import json
from com.demo1.ai.kimi import demo1 as textAi
from com.demo1.ai.kimi import file_demo as fileAi
from com.demo1.ai.qianfan import chat_demo as qfTextAi
from PIL import Image, ImageTk


def append_text(text_widget, text_content):
    # 添加文本内容
    text_widget.configure(state='normal')  # 临时切换到normal状态以插入文本
    text_widget.insert(tk.END, text_content)
    text_widget.see(tk.END)
    text_widget.configure(state='disabled')  # 再次设置为只读状态
    text_widget.pack(expand=True, fill=tk.BOTH)


app = tk.Tk()
app.title("AI小工具")

# 设置窗口大小
app.geometry("600x800")

"""
    创建一个只用于展示的Text widget
    :param master: 父容器
    """
text_widget = tk.Text(app, wrap=tk.WORD, state='disabled', bg='white', relief=tk.FLAT)

# 添加文本内容
text_content = "AI：你好，这是一个简单的AI对话窗口，你可以输入任何你想问的问题，我给你回答。\n-------------------------------------------------\n"
text_widget.configure(state='normal')  # 临时切换到normal状态以插入文本
text_widget.insert(tk.END, text_content)
text_widget.configure(state='disabled')  # 再次设置为只读状态
text_widget.pack(expand=True, fill=tk.BOTH)

text_area = tk.Text(app, width=50, height=20)
text_area.pack(side=tk.LEFT)

image_name = None
file_name = None


def read_text_and_append(channel):
    global image_name
    global file_name
    user_msg = text_area.get("1.0", tk.END)
    user_msg = user_msg.strip()
    if channel == "KM_NO_STREAM":
        t = Thread(target=ai_msg_no_stream, args=(user_msg,))
        t.start()
    elif channel == "KM_STREAM":
        t = Thread(target=ai_msg_stream, args=(user_msg,))
        t.start()
    elif channel == "QF_NO_STREAM":
        t = Thread(target=qf_ai_msg_no_stream, args=(user_msg,))
        t.start()
    elif channel == "QF_STREAM":
        t = Thread(target=qf_ai_msg_stream, args=(user_msg,))
        t.start()
    elif channel == "READ_FILE":
        t = Thread(target=ai_read_file, args=(user_msg,))
        t.start()
        file_name = None
    elif channel == "READ_IMG":
        replace_text = f"图片：{image_name}"
        user_msg = user_msg.replace(replace_text, "")
        t = Thread(target=ai_read_image, args=(user_msg,))
        t.start()
        append_text(text_widget, "\n用户：")
        print(f"filename:{image_name}")
        image = Image.open(image_name)
        image_resized = image.resize((200, 100))
        imagetk = ImageTk.PhotoImage(image_resized)
        text_widget.insert(tk.END, "\n")
        text_widget.image_create(tk.END, image=imagetk)
        text_widget.image = imagetk
        image_name = None
    append_text(text_widget, "\n用户：{}".format(user_msg))
    text_area.delete("1.0", tk.END)


def choose_image():
    global image_name
    image_name = filedialog.askopenfilename(filetypes=[("Image Files", "*.png *.jpg *.jpeg")])
    print(f"filename:{image_name}")
    if image_name:
        # 放入输入框中
        text_area.insert(tk.END, f"\n 图片：{image_name}\n")


def ai_msg_no_stream(user_msg):
    ai_text = textAi.no_stream_v1(user_msg)
    append_text(text_widget, "\nAI:{}".format(ai_text))


def ai_msg_stream(user_msg):
    append_text(text_widget, "\nAI：")
    response = textAi.stream_v1(user_msg)
    for idx, chunk in enumerate(response):
        # print("Chunk received, value: ", chunk)
        chunk_message = chunk.choices[0].delta
        if not chunk_message.content:
            continue
        append_text(text_widget, chunk_message.content)


def qf_ai_msg_no_stream(user_msg):
    ai_text = qfTextAi.no_stream_v1(user_msg)
    append_text(text_widget, "\nAI:{}".format(ai_text))


def ai_read_image(user_msg):
    ai_text = fileAi.ai_msg_demo(image_name, user_msg)
    append_text(text_widget, "\nAI:{}".format(ai_text))


def ai_read_file(user_msg):
    ai_text = fileAi.ai_file_demo(file_name, user_msg)
    append_text(text_widget, "\nAI:{}".format(ai_text))


def qf_ai_msg_stream(user_msg):
    append_text(text_widget, "\nAI：")
    response = qfTextAi.stream_v1(user_msg)
    for line in response.iter_lines():
        if line:
            append_text(text_widget, json.loads(line.decode("UTF-8").split("data: ")[1])["result"])


button = tk.Button(app, text="kimi", width=4, height=2, command=lambda: read_text_and_append("KM_NO_STREAM"))
button.pack()
button = tk.Button(app, text="kimi（流式）", width=8, height=2, command=lambda: read_text_and_append("KM_STREAM"))
button.pack()

button = tk.Button(app, text="千帆", width=4, height=2, command=lambda: read_text_and_append("QF_NO_STREAM"))
button.pack()
button = tk.Button(app, text="千帆（流式）", width=8, height=2, command=lambda: read_text_and_append("QF_STREAM"))
button.pack()

button = tk.Button(app, text="选择图片", width=8, height=2, command=choose_image)
button.pack()
button = tk.Button(app, text="解读图片", width=8, height=2, command=lambda: read_text_and_append("READ_IMG"))
button.pack()
button = tk.Button(app, text="解读文件", width=8, height=2, command=lambda: read_text_and_append("READ_FILE"))
button.pack()

app.mainloop()
